package rootkit

import (
	"syscall"
	"os"
	"fmt"
	"io/ioutil"
	"time"
	"strconv"
	"cloud_worm/util/uexec"
)

const (
	NoSuchProcesErr = "no such process"
)

// pid的rootkit的检查主要体现在以下几个方面：
// 		1. 异常数超过15条，则停止检查，直接报错；
//		2. 异常进程的检查

type PidCheck struct {
	NoProc	 		bool
	Ps 				string

	Stat 			StatInfo				// 统计信息
}

func (p *PidCheck) CheckPid() error {
	// 确定ps命令
	psCmd := "/bin/ps"
	if _, err := os.Stat(psCmd); err != nil {
		psCmd = "/usr/bin/ps"
		if _, err := os.Stat(psCmd); err != nil {
			psCmd = ""
		}
	}
	p.Ps = psCmd

	// 确定/proc目录的存在
	proc0 := "/proc"
	proc1 := "/proc/1"
	_, err1 := os.Stat(proc0)
	_, err2 := os.Stat(proc1)
	if err1 == nil && err2 == nil {
		p.NoProc = false
	} else {
		p.NoProc = true
	}

	// 暂定max pid值为：32768
	p.scanAllPids(32768)

	// TODO: 告警个数的处理

	return nil
}

func (p *PidCheck) scanAllPids(maxPid int) {
	myPid := os.Getpid()

	for i := 1; i <= maxPid; i++ {
		p.Stat.Total++

		// 忽略本身进程pid
		if i == myPid {
			continue
		}

		// process pid test
		kill0 := false
		sid0 := false
		gid0 := false
		if err := syscall.Kill(i, 0); err != nil && err.Error() == NoSuchProcesErr {
			kill0 = true
		}
		if _, err := syscall.Getsid(i); err != nil && err.Error() == NoSuchProcesErr {
			sid0 = true
		}
		if _, err := syscall.Getpgid(i); err != nil && err.Error() == NoSuchProcesErr {
			gid0 = true
		}
		// proc pid test
		stat, read, chdir := p.procCheck(i)

		// 确实确认进程不存在，则跳过
		if kill0 == false && sid0 == false && gid0 == false &&  stat == false && read == false && chdir == false {
			continue
		}

		if p.Stat.Errors > 15 {
			// TODO: 错误数达到一定值就开始报警可以了，不用全量扫描完，暂定15条
			fmt.Sprintf("Excessive number of hidden processes. It maybe a false-positive or something really bad is going on.")
			return
		}

		// 尝试执行ps获取进程信息
		ps0 := -1
		if p.Ps != "" {
			ps0 = 0
			psCmd := fmt.Sprintf("%s -p %d > /dev/null 2>&1", p.Ps, i)
			msg, ret := uexec.ExecUnixCmd(psCmd, 60)
			if ret == 0 && msg == "" {
				ps0 = 1
			}
		}
		// sleep 避免造成冲击
		time.Sleep(2 * time.Second)

		// 说明进程确实存在，则不用管
		if ps0 != 0 && kill0 && sid0 && gid0 && stat && read {
			continue
		}

		// 如果通过接口getsid、kill等获取到pid，但是ps(1)没有出现，则需要找出该pid是否被删除，或者被隐藏
		kill1 := false
		sid1 := false
		gid1 := false
		if err := syscall.Kill(i, 0); err != nil && err.Error() == NoSuchProcesErr {
			kill1 = true
		}
		if _, err := syscall.Getsid(i); err != nil && err.Error() == NoSuchProcesErr {
			sid1 = true
		}
		if _, err := syscall.Getpgid(i); err != nil && err.Error() == NoSuchProcesErr {
			gid1 = true
		}
		stat, read, chdir = p.procCheck(i)
		// 2s后的再次查，发现进程已经没有了，则截止
		if kill1 == false && sid1 == false && gid1 == false && stat == false && read == false && chdir == false {
			continue
		}

		// 异常诊断
		if sid0 == sid1 && kill0 == kill1 && sid0 != kill0 {
			// kill发现有，但是getsid和getpgid却没有发现，可能是一个缺省进程，忽略
			if !(kill0 == true && sid0 == false && gid0 == false && sid1 == false) {
				// TODO: 告警
				fmt.Sprintf("process '%d' hidden from kill(%v) or getsid(%v). Possible kernel-level rootkit", i, kill0, sid0)
				p.Stat.Errors++
			}
		} else if kill1 != sid1 || gid1 != kill1 || gid1 != sid1 {
			if !(kill1 == true && sid1 == false && gid0 == false) {
				// TODO：告警
				fmt.Sprintf("process '%d' hidden from kill(%v) or getsid(%v). Possible kernel-level rootkit", i, kill1, sid1)
				p.Stat.Errors++
			}
		} else if read != stat || read != chdir || stat != kill1 {
			// 如果pid是线程thread，则不会有/proc中对应的目录，需要进一步检查
			if p.NoProc == false && p.checkForProcMayThread(i) == false {
				// TODO: 告警
				fmt.Sprintf("Process '%d' hidden from /proc. Possible kernel level rootkit", i)
				p.Stat.Errors++
			}
		} else if sid1 && kill1 && ps0 == 0 {
			// 如果一个pid是一个线程thread，则不会在ps中展示，需要进一步检查
			if p.checkForProcMayThread(i) == false {
				// TODO: 告警
				fmt.Sprintf("Process '%d' hidden from ps. Possible trojaned version installed.", i)
				p.Stat.Errors++
			}
		}
	}
	return
}

func (p *PidCheck) procCheck(pid int) (stat, read, chdir bool) {
	stat = false
	read = false
	chdir = false
	if p.NoProc {
		return
	}

	filePath := fmt.Sprintf("/proc/%d", pid)
	_, err := os.Stat(filePath)
	if err != nil {
		stat = false
	} else {
		stat = true
	}

	read = p.isFileOnProcDir(pid)
	chdir = p.dealChDir(pid)

	return
}

func (p *PidCheck) isFileOnProcDir(pid int) bool {
	fileInfos, err := ioutil.ReadDir("/proc")
	if err != nil {
		return false
	}

	pidFile := fmt.Sprintf("%d", pid)
	for _, item := range fileInfos {
		if item.Name() == pidFile {
			return true
		}
	}
	return false
}

func (p *PidCheck) dealChDir(pid int) bool {
	currDir, err := os.Getwd()
	if err != nil {
		return false
	}
	defer os.Chdir(currDir)

	if err := os.Chdir("/proc"); err != nil {
		return false
	}
	if err := os.Chdir(fmt.Sprintf("/proc/%d", pid)); err != nil {
		return false
	}

	return true
}

func (p *PidCheck) checkForProcMayThread(pid int) bool {
	if _, err := os.Stat(fmt.Sprintf("/proc/.%d", pid)); err == nil {
		return true
	}

	findThreadPid := &FindThreadPid{ProFoundPid: false}
	findThreadPid.readProcDir("/proc", fmt.Sprintf("%d", pid), Proc)
	return findThreadPid.ProFoundPid
}

const (
	Proc = iota
	Pid
	Task
)

type FindThreadPid struct {
	ProFoundPid 		bool
}

func (f *FindThreadPid) readProcDir(dirName string, pid string, position int) {
	fileInfos, err := ioutil.ReadDir(dirName)
	if err != nil {
		return
	}

	for _, item := range fileInfos {
		if f.ProFoundPid {
			break
		}

		if item.Name() == "." || item.Name() == ".." {
			continue
		}

		if position == Proc {
			// 当前位置为: /proc
			if _, err := strconv.Atoi(item.Name()); err != nil {
				continue
			}

			fileName := fmt.Sprintf("%s/%s", dirName, item.Name())
			f.readProcFile(fileName, pid, position + 1)
		} else if position == Pid {
			// 当前位置为: /proc/pid
			if item.Name() == "task" {
				fileName := fmt.Sprintf("%s/%s", dirName, item.Name())
				f.readProcFile(fileName, pid, position + 1)
			}
		} else if position == Task {
			// 当前位置为: /proc/pid/task
			if item.Name() == pid {
				f.ProFoundPid = true
				break
			}
		} else {
			break
		}
	}

	return
}

func (f *FindThreadPid) readProcFile(fileName string, pid string, position int) {
	fInfo, err := os.Stat(fileName)
	if err != nil {
		return
	}

	if fInfo.IsDir() {
		f.readProcDir(fileName, pid, position)
		return
	}

	return
}